/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

syntax = "proto3";

package tensorflow;

import "tensorflow/compiler/xla/service/hlo.proto";
import "tensorflow/compiler/xla/xla_data.proto";
import "tensorflow/core/framework/tensor_shape.proto";

// A serialization of TPUExecutable. Only includes fields necessary to load
// and execute a program on a worker node.
message TPUExecutableInfoProto {
  reserved 1;

  // The shapes of the inputs and outputs.
  repeated xla.ShapeProto input_shapes = 2;
  reserved 7;  // was input_shape
  xla.ShapeProto output_shape = 3;

  message UpdateIndexPair {
    int32 index = 1;
    bool updated = 2;
  }

  message ShapeIndex {
    repeated int32 index = 1;
  }

  // Dynamic output indices indicate which outputs have dynamic dimensions.
  repeated ShapeIndex dynamic_output_indices = 11;

  // For each resource variable output, what was the index of the corresponding
  // input and was it updated? The indices are sorted by input order.
  repeated UpdateIndexPair variable_indices = 10;

  // The shapes of the outputs when represented as Tensors. These may not
  // match the output_shape values because we may flatten tensors to avoid
  // excess padding.
  repeated TensorShapeProto output_tensor_shapes = 8;

  reserved 4;

  // Optional session module for passing XLA computations between TPUCompileOp
  // and TPUExecuteOp. This is needed to support the
  // --xla_dump_hlo_snapshots flag.
  xla.HloSnapshot session_module = 5;

  // The physical device ids assigned to the replicated cores.
  xla.DeviceAssignmentProto device_assignment = 6;
}

message TPUHostTransferInfoProto {
  reserved 1;
  // Serialized metadata for a data transfer between device and host.
  repeated bytes serialized_transfers = 2 [ctype = CORD];
}
